Greptile Summary

This PR updates the cuDNN-based fused attention forward pass to always return Stats (= log(SumExp) + Max) directly from cuDNN, removing the prior manual Stats = log(SumExp) + Max computation in Python. When return_max_logit=True, cuDNN also returns the Max tensor. The Sum_Exp output is eliminated entirely. The descriptor field generate_max_sum_exp is renamed to return_max_logit, and the aux tensor ordering is updated throughout (Stats is now always index 0; Max is index 1 when present).

Key changes:

• fused_attn_f16_arbitrary_seqlen.cu: generate_stats is hardcoded to true; Sum_Exp tensor removed; Stats is always set as output; aux-tensor indexing restructured so Stats is always first.
• fused_attn.py: Removes stats = output_tensors[1] + torch.log(output_tensors[2]) computation; reads Stats and Max directly from cuDNN output at indices 1 and 2.
• utils.h: FADescriptor_v1::generate_max_sum_exp renamed to return_max_logit.
• fused_attn.h: Two of the four return_max_logit doc comments updated (per prior review threads, the other two at lines 403 and 556 still need updates).

Two SM120 (Blackwell) regressions were found in fused_attn_f16_arbitrary_seqlen.cu:

1. The Stats stride-setting condition uses is_ragged_q && cudnn_runtime_version >= 90600 instead of use_ragged_stats, bypassing the sm_arch_ != 120 guard and calling set_ragged_offset(nullptr) on SM120.
2. The output_S allocation shape condition also drops the sm_arch_ != 120 guard, causing a 3D/4D shape mismatch on SM120 with THD format.

Confidence Score: 2/5

• Two SM120-specific regressions in the C++ CUDA kernel should be fixed before merging.
• The overall design is sound and the non-SM120 path looks correct. However, two concrete bugs on SM120 (Blackwell) were introduced: (1) the Stats stride condition drops the sm_arch_ != 120 guard present in use_ragged_stats, passing a null offset_stats pointer to set_ragged_offset() during cuDNN graph construction; (2) the output_S shape allocation also drops the SM120 guard, causing a 3D/4D shape mismatch with what the graph actually builds on SM120. Since Blackwell is a supported production architecture, these are real correctness issues that should be addressed before merging.
• transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu — both the graph-building Stats stride condition (~line 384) and the output_S shape allocation (~line 1134) need the sm_arch_ != 120 guard restored.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[fused_attn_arbitrary_seqlen_fwd] --> B{Aux_CTX_Tensors empty?}
    B -->|Yes - allocate| C[Allocate Stats tensor\nalways, index 0]
    C --> D{return_max_logit?}
    D -->|Yes| E[Allocate Max tensor\nindex 1]
    D -->|No| F[Skip Max]
    E --> G[Allocate rng_state\nindex 2]
    F --> G2[Allocate rng_state\nindex 1]
    B -->|No - use existing| H[Read Stats → devPtrS1\nindex 0]
    H --> I{return_max_logit?}
    I -->|Yes| J[Read Max → devPtrS2\nindex 1]
    I -->|No| K[Skip Max]
    J --> L[Read rng_state\nindex 2]
    K --> L2[Read rng_state\nindex 1]

    subgraph graph_builder [Graph Builder - fused_attn_arbitrary_seqlen_fwd_impl]
        M[sdpa generates O, Stats always] --> N[Stats: set_output=true\nset stride always]
        N --> O{return_max_logit?}
        O -->|Yes| P[Max tensor\nset_logit_max]
        O -->|No| Q[Stats only\nStats_tuple = Stats, null]
        P --> R[Stats_tuple = Stats, Max]
    end

    subgraph python [Python fused_attn_fwd - return_max_logit=True]
        S[output_tensors: out, Stats, Max, rng_state, ...]
        S --> T[aux_ctx_tensors = Stats + rng_state + optional]
        S --> U[max_tensor = output_tensors 2 = Max]
        U --> V[max_logit = amax over batch/seq dims]
        T --> W[return out, aux_ctx_tensors, max_logit]
    end

_{Reviews (19): Last reviewed commit: "Merge branch 'fix_return_stats_max_cudnn..." | Re-trigger Greptile}

/te-ci L2